#coding = utf-8

'''
生成训练数据
'''

import os
import SimpleITK as sitk
import numpy as np
from PIL import Image
from skimage import measure
import cv2
import matplotlib.pyplot as plt
import joblib
from pathlib2 import Path

def garbor_filter4(image):
    image = image * 255
    image = image.astype(np.uint8)


    filters3 = []
    result3_list = []
    #theta2 = [0, 1 * np.pi / 6, 2 * np.pi / 6, 3 * np.pi / 6, 4 * np.pi / 6, 5 * np.pi / 6]
    theta2 = [0, 1 * np.pi / 6, 2 * np.pi / 6, 3 * np.pi / 6, 4 * np.pi / 6, 5 * np.pi/6]
    for item in theta2:
        kern = cv2.getGaborKernel((2, 2), sigma=1.0, theta=item, lambd=np.pi / 2.0, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters3.append(kern)
    for i in range(len(filters3)):
        accum = np.zeros_like(image)
        for kern in filters3[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result3_list.append(np.array(accum))


    return result3_list

#将数据保存到chengkun格式
def transfer_data_to_chengkun():
    predict_root_path = "E:\predict\gabor_class_instance8"
    raw_root_path = "E:\Dataset\Qiye"
    save_root_path = "F:\Dataset\chengkun_liver_cluster3"

    model = joblib.load("kmeans_3.model")

    raw_id_list = []
    for item in sorted(os.listdir(os.path.join(raw_root_path, "DongBeiDaXue\liver"))):
        raw_id_list.append(item.split("_")[1])
    for item in sorted(os.listdir(os.path.join(raw_root_path, "DongBeiDaXue2\liver"))):
        raw_id_list.append(item.split("_")[1])

    for i in range(80):
        case_id = "case_{}".format(str(i).zfill(5))

        if i < 50:
            raw_liver_file = os.path.join(raw_root_path, "DongBeiDaXue\liver\\data2_{}_liver_label.mha".format(raw_id_list[i]))
            raw_tumor_file = os.path.join(raw_root_path, "DongBeiDaXue\lesion\\data2_{}_lesion_label.mha".format(raw_id_list[i]))
            raw_image_file = os.path.join(raw_root_path, "DongBeiDaXue\image_venous\\data2_{}_venous.mha".format(raw_id_list[i]))
        else:
            raw_liver_file = os.path.join(raw_root_path, "DongBeiDaXue2\liver\\data2_{}_liver_label.mha".format(raw_id_list[i]))
            raw_tumor_file = os.path.join(raw_root_path,
                                          "DongBeiDaXue2\lesion\\data2_{}_lesion_label.mha".format(raw_id_list[i]))
            raw_image_file = os.path.join(raw_root_path, "DongBeiDaXue2\image_venous\\data2_{}_venous.mha".format(
                raw_id_list[i]))

        predict_tumor_path = os.path.join(predict_root_path, "{}\predict_tumor".format(case_id))
        predict_prob_path = os.path.join(predict_root_path, "{}\prob_numpy".format(case_id))

        raw_liver = sitk.GetArrayFromImage(sitk.ReadImage(raw_liver_file))
        raw_tumor = sitk.GetArrayFromImage(sitk.ReadImage(raw_tumor_file))
        raw_image = sitk.GetArrayFromImage(sitk.ReadImage(raw_image_file))

        liver_list = []
        tumor_list = []
        image_list = []
        for t in range(raw_liver.shape[0]):
             if np.max(raw_liver[t]) > 0:
                 liver_list.append(raw_liver[t])
                 tumor_list.append(raw_tumor[t])
                 image_list.append(raw_image[t])
        raw_tumor = np.array(tumor_list)
        raw_liver = np.array(liver_list)
        raw_image = np.array(image_list)

        #获取预测结果
        predict_tumor = np.zeros(raw_tumor.shape)
        predict_prob = np.zeros(raw_tumor.shape)

        for item in sorted(os.listdir(predict_tumor_path)):
            file_name = os.path.join(predict_tumor_path, item)
            index = int(item.strip().split(".")[0])
            data = Image.open(file_name).convert("L")
            data = np.array(data)
            data = data / 255
            data[data>0] = 1
            predict_tumor[index] = data
        for item in sorted(os.listdir(predict_prob_path)):
            file_name = os.path.join(predict_prob_path, item)
            index = int(item.strip().split(".")[0])
            data = np.load(file_name)
            predict_prob[index] = data

        dice = 2 * (predict_tumor*raw_tumor).sum() / (predict_tumor.sum() + raw_tumor.sum())
        print(case_id, dice)

        #获取得到图像的滤波结果
        raw_image[raw_image >= 250] = 250
        raw_image[raw_image <= -200] = -200
        raw_image = (raw_image + 200) /450
        gabor_result0 = []
        gabor_result1 = []
        gabor_result2 = []
        gabor_result3 = []
        for t in range(raw_image.shape[0]):
            image_temp = raw_image[t] * raw_liver[t]
            gabor_result = garbor_filter4(image_temp)
            gabor_result0.append(gabor_result[0])
            gabor_result1.append(gabor_result[1])
            gabor_result2.append(gabor_result[2])
            gabor_result3.append(gabor_result[3])
        gabor_result0 = np.array(gabor_result0)
        gabor_result1 = np.array(gabor_result1)
        gabor_result2 = np.array(gabor_result2)
        gabor_result3 = np.array(gabor_result3)
        gabor_result0 = gabor_result0 / 255.
        gabor_result1 = gabor_result1 / 255.
        gabor_result2 = gabor_result2 / 255.
        gabor_result3 = gabor_result3 / 255.




        #获取肿瘤实例
        [tumor_labels, num] = measure.label(raw_tumor, return_num=True)
        region = measure.regionprops(tumor_labels)
        for t in range(num):
           avg_slice = int(region[t].area / (region[t].bbox[3] - region[t].bbox[0]))
           if avg_slice < 10 :
               continue
           label_temp = np.zeros(tumor_labels.shape)
           label_temp[tumor_labels == t+1] = 1
           recall = float((label_temp[predict_tumor == 1] == 1).sum()) / float((label_temp == 1).sum())

           # 计算相关的特征
           gabor0 = (label_temp * gabor_result0).sum() / label_temp.sum()
           gabor1 = (label_temp * gabor_result1).sum() / label_temp.sum()
           gabor2 = (label_temp * gabor_result2).sum() / label_temp.sum()
           gabor3 = (label_temp * gabor_result3).sum() / label_temp.sum()

           predict = (label_temp * predict_prob).sum() / label_temp.sum()

           feature = [gabor0, gabor1, gabor2, gabor3, predict, recall]
           feature = np.array(feature)
           feature = feature.reshape((1,feature.shape[0]))
           label = model.predict(feature)[0]

           #设置标签
           raw_liver[label_temp == 1] = 2 + label

        print(np.unique(raw_liver))
        raw_liver = raw_liver.astype(np.uint8)
        raw_image = raw_image.astype(np.float32)

        destination_patient_path = os.path.join(save_root_path, case_id)
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        for j in range(raw_image.shape[0]):
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), raw_image[j])
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), raw_liver[j])

        print()











if __name__ == '__main__':
    transfer_data_to_chengkun()